#-*- coding: utf-8 -*-
import numpy as np
import math
import pat.utils     as mu
import pat.algebra   as ma

#
# Some indexes for using np arrays as state and input vectors
#
# state vector
s_x     = 0   # position x axis
s_y     = 1   # position y axis
s_z     = 2   # heigh above ground
s_v     = 3   # airspeed
s_alpha = 4   # alpha
s_beta  = 5   # beta
s_phi   = 6   # roll  (euler, ltp to body)
s_theta = 7   # pitch (euler, ltp to body)
s_psi   = 8   # yaw   (euler, ltp to body)
s_p     = 9   # rotational vel body x
s_q     = 10  # rotational vel body y
s_r     = 11  # rotational vel body z
s_size  = 12

# input vector
i_th   = 0     # throttle
i_da   = 1     # aileron
i_de   = 2     # elevator
i_dr   = 3     # rudder
i_size = 4

# tmp hack for introducing multiple surfaces and multiple engines
i_da1 = 0
i_de1 = 1
i_dr1 = 2

def get_rho(h):
    """
    Get air density at the given altitude (in m)
    Warning, the two parts of that function do not
    join not well!
    """
    if h<= 11000:
        return 1.225 * math.pow((1-6.5e-3*h/288.15), 4.2557)
    else:
        return 0.36 * math.exp(-1.17e-4*(h-11000))

def get_aero_to_body(X):
    """
    computes the aero to body rotation matix
    """
    ca = math.cos(X[s_alpha]); sa = math.sin(X[s_alpha]) 
    cb = math.cos(X[s_beta]);  sb = math.sin(X[s_beta])
    return np.array([[ca*cb, -ca*sb, -sa],
                     [sb   ,  cb   ,  0.],
                     [sa*cb, -sa*sb,  ca]])

def get_f_eng_body(X, U, P):
    """
    return propulsion forces expressed in body frame
    """
    rho = get_rho(-X[s_z])
    f_engines_body = np.zeros((P.eng_nb, 3))
    for i in range(0, P.eng_nb):
        thrust = U[i]*P.fmaxs[i]*math.pow((rho/P.rhois[i]),P.nrhos[i])*math.pow((X[s_v]/P.Vis[i]),P.nVs[i]) 
        f_engines_body[i] = np.dot(P.eng_to_body[i], np.array([thrust, 0., 0.]))
    return f_engines_body 

def get_f_aero_body(X, Usfc, P, Pdyn):
    """
    return aerodynamic forces expressed in body frame
    """
    d_alpha = X[s_alpha] - P.alpha0
    X_rvel_body = X[s_p:s_r+1]

    CL = P.CL0 + P.CL_alpha*d_alpha + P.CL_beta*X[s_beta] +\
         np.dot(P.CL_omega,X_rvel_body) + np.dot(P.CL_sfc,Usfc)

    CD = P.CD0 + P.CD_k1*CL + P.CD_k2*(CL**2) + np.dot(P.CD_sfc,Usfc)

    CY = P.CY_alpha*d_alpha + P.CY_beta*X[s_beta] +\
         np.dot(P.CY_omega,X_rvel_body) + np.dot(P.CY_sfc,Usfc)
    
    return Pdyn*P.Sref*np.dot(get_aero_to_body(X),[-CD, CY, -CL])

def get_m_eng_body(f_eng_body, P):
    """
    return propulsion moments expressed in body frame
    """
    m = np.zeros(3)
    for i in range(0, P.eng_nb):
        m += np.cross(P.eng_pos[i], f_eng_body[i])
    return m

def get_m_aero_body(X, Usfc, P, Pdyn):
    """
    return aerodynamic moments expressed in body frame
    """
    d_alpha = X[s_alpha] - P.alpha0
    X_rvel_body = X[s_p:s_r+1]

    Cl = P.Cl_alpha*d_alpha + P.Cl_beta*X[s_beta] +\
         np.dot(P.Cl_omega,X_rvel_body) + np.dot(P.Cl_sfc,Usfc)
    Cm = P.Cm_alpha*d_alpha + P.Cm_beta*X[s_beta] +\
         np.dot(P.Cm_omega,X_rvel_body) + np.dot(P.Cm_sfc,Usfc)
    Cn = P.Cn_alpha*d_alpha + P.Cn_beta*X[s_beta] +\
         np.dot(P.Cn_omega,X_rvel_body) + np.dot(P.Cn_sfc,Usfc)

    return Pdyn*P.Sref*np.array([Cl*P.Bref, Cm*P.Cref, Cn*P.Bref])

def dyn(X, t, U, P):
    """
    Dynamic model
    """
    rho = get_rho(-X[s_z])
    Pdyn = 0.5*rho*X[s_v]**2

    Ueng = U[0:P.eng_nb]                  # engines part of input vector
    Usfc = U[P.eng_nb:P.eng_nb+P.sfc_nb]  # control surfaces part of input vector
    d_alpha = X[s_alpha] - P.alpha0       # 
    X_rvel_body = X[s_p:s_r+1]            # body rotational velocities
    X_euler = X[s_phi:s_psi+1]            # euler angles                          

    # Newton for forces in body frame
    f_aero_body = get_f_aero_body(X, Usfc, P, Pdyn)
    f_eng_body  = get_f_eng_body(X, Ueng, P)
    earth_to_body = ma.rmat_of_euler(X_euler)
    f_weight_body = np.dot(earth_to_body, [0., 0., P.m*P.g])
    forces_body = f_aero_body + np.sum(f_eng_body, axis=0) + f_weight_body
    
    vel_body = np.dot(get_aero_to_body(X), [X[s_v], 0., 0.])  # u, v, w
    accel_body = 1./P.m*forces_body - np.cross(X_rvel_body, vel_body)

    # Newton for moments in body frame
    m_aero_body = get_m_aero_body(X, Usfc, P, Pdyn)
    m_eng_body = get_m_eng_body(f_eng_body, P)
    raccel_body = np.dot(P.invI, m_aero_body + m_eng_body - np.cross(X_rvel_body, np.dot(P.I, X_rvel_body)))
    
    Xdot = np.zeros(s_size)
    Xdot[s_x:s_z+1] = np.dot(np.transpose(earth_to_body), vel_body)

    Xdot[s_v] = np.inner(vel_body, accel_body)/X[s_v]
    u, v, w = vel_body
    ud, vd, wd = accel_body
    Xdot[s_alpha] = (u*wd - w*ud)/(u**2+w**2)
    Xdot[s_beta] = (X[s_v]*vd - v*Xdot[s_v]) / X[s_v] / math.sqrt(u**2+w**2)

    Xdot[s_phi:s_psi+1] = ma.euler_derivatives(X_euler, X_rvel_body)
    
    Xdot[s_p:s_r+1] = raccel_body

    return Xdot


#
#
# From here we use Objects
#
#

from scipy import integrate
import gobject
import StringIO
import xml.etree.ElementTree as ET

class Param:
    """
    Represents the parameters of a flight dynamic model
    """
    def __init__(self, filename=None):
        self.g = 9.81
        self.read_from_xml(filename)
        self.compute_auxiliary()

    def compute_auxiliary(self):
        """
        compute auxiliary coefficients
        """
        self.invI = np.linalg.inv(self.I)
        self.eng_to_body=[]
        for i in range(0, self.eng_nb):
            self.eng_to_body.append(np.eye(3))
        self.input_nb = self.eng_nb+self.sfc_nb
            

    def read_from_xml(self, filename="../config/Navion.xml"):
        """
          Quick and dirty xml parsing
        """
        print "Info: loading parameters from: {:s}".format(filename)
        tree = ET.parse(filename)
        root = tree.getroot()

        self.name = root.attrib["name"]

        masses = root.find("masses")
        self.m = float(masses.find("mass").text)
        self.I = np.genfromtxt(StringIO.StringIO(masses.find("inertia").text), delimiter=",")

        aerodynamics = root.find("aerodynamics")
        for param in ["Sref", "Cref", "Bref"]:
            setattr(self, param, float(aerodynamics.find(param).text))

        stability = aerodynamics.find("stability")
        misc = stability.find("misc")
        for att in ['CL0', 'alpha0', 'CD0', 'CD_k1', 'CD_k2']:
            setattr(self, att, float(misc.attrib[att]))
        for s in ['alpha', 'beta', 'p', 'q', 'r']:
            x = stability.find(s)
            for c in ["CL", "CY", "Cl", "Cm", "Cn"]:
                setattr(self, c+'_'+s, float(x.attrib[c]))
        for c in ["CL", "CY", "Cl", "Cm", "Cn"]: # CL_omega
            val = [float(stability.find(s).attrib[c]) for s in ['p', 'q', 'r']]
            setattr(self, c+'_omega', val)
            
        surfaces = stability.findall("surface");
        self.sfc_name = [s.attrib["name"] for s in surfaces]
        self.sfc_nb = len(self.sfc_name)
        for c in ["CL", "CY", "CD", "Cl", "Cm", "Cn"]:
            setattr(self, c+"_sfc", np.array([float(s.attrib[c]) for s in surfaces]))

        propulsion = root.find("propulsion")
        engines = propulsion.findall("engine");
        self.eng_name = [e.attrib["name"] for e in engines]
        self.eng_nb = len(self.eng_name)
        setattr(self, "eng_pos", [np.genfromtxt(StringIO.StringIO(e.attrib["pos"]),delimiter=",") for e in engines])
        setattr(self, "eng_align", [np.genfromtxt(StringIO.StringIO(e.attrib["align"]),delimiter=",") for e in engines])
        for param in ["fmax", "rhoi", "nrho", "Vi", "nV", "tau"]:
            setattr(self, param+"s", [float(e.attrib[param]) for e in engines])


    def __str__(self):
        """
          Quick and dirty human readable printing of parameters
        """
        res  = "Name: {:s}".format(self.name)
        res += """
  Mass: {:f} kg
  Inertia:\n {:s} kg.m2
        """.format(self.m, str(self.I))
        res += """
  Aero:
    Sref: {:f} m2    
    Cref: {:f} m    
    Bref: {:f} m    
    """.format(self.Sref, self.Cref, self.Bref)
        for m in ['CL0', 'alpha0', 'CD0', 'CD_k1', 'CD_k2']: res += (m+": {:f} ".format(getattr(self, m)))
        res+="\n           alpha      beta         p         q         r\n"
        for c in ["CL", "CY", "Cl", "Cm", "Cn"]:
            res+= "    {:s}".format(c)
            for s in ['alpha', 'beta', 'p', 'q', 'r']:
                res+=" {: f}".format(getattr(self, c+"_"+s))
            res+="\n"
        res += "      "
        for name in self.sfc_name: res+= "{:>10s}".format(name)
        res+="\n"
        for c in ["CL", "CY", "CD", "Cl", "Cm", "Cn"]:
            res+= "    {:s}".format(c)
            for i in range(0, len(self.sfc_name)):
                res+=" {: f}".format(getattr(self, c+"_sfc")[i])
            res+="\n"

        res += "\n  Propulsion:\n"
        res += "          "
        for name in self.eng_name: res+= "{:>16s}".format(name)
        res+="\n"

        for param in ["pos", "align"]:
            res+="    {:8s}".format(param)
            for i in range(0, self.eng_nb):
                res += "{:>16s}".format(getattr(self, "eng_"+param)[i])
            res+="\n"

        for param in ["fmax", "rhoi", "nrho", "Vi", "nV", "tau"]:
            res += "    {:8s}".format(param)
            for i in range(0, self.eng_nb):
                res += "{:16f}".format(getattr(self, param+"s")[i])
            res+="\n"
        return res
        
class State():
    def __init__(self, x=0., y=0., z=0., v=68., alpha=0., beta=0., phi=0., theta=0., psi=0., p=0., q=0., r=0.):
        self.x = x; self.y = y; self.z = z
        self.v = v; self.alpha = alpha; self.beta=beta
        self.phi = phi; self.theta = theta; self.psi = psi
        self.p = p; self.q = q; self.r = r

    def as_array(self): return np.array([self.x, self.y, self.z,
                                         self.v, self.alpha, self.beta,
                                         self.phi, self.theta, self.psi,
                                         self.p, self.q, self.r])

class DynamicModel(gobject.GObject):
    __gsignals__ = { 
    "update_fast" : (gobject.SIGNAL_RUN_LAST, gobject.TYPE_NONE,()),
    "update_slow" : (gobject.SIGNAL_RUN_LAST, gobject.TYPE_NONE,())
    } 

    sta_running, sta_crashed = range(2)

    def __init__(self, filename="../config/Rcam_single_engine.xml", X0=None, U0=None):
        gobject.GObject.__init__(self)
        self.P = Param(filename)
        if X0 <> None: self.X0 = X0
        else: self.X0 = State(theta = mu.rad_of_deg(5.)).as_array()
        self.U = np.zeros(self.P.input_nb)
        self.reset()
#        from pdb import set_trace; set_trace()
        try:
            self.c_impl = __import__("dynamic_model_c")
            self.c_impl.init()
        except ImportError:
            self.c_impl = None
            print "Info: could not load dynamic model C library, defaulting to pure Python"

    def reset(self):
        self.X = np.array(self.X0)
        self.emit("update_slow")
        self.emit("update_fast")
        self.update_cnt = 0
        self.status = DynamicModel.sta_running

    def run(self, U, dt):
        if self.status == DynamicModel.sta_crashed: return
        self.U = U
        
        if self.c_impl == None:
            foo, self.X = integrate.odeint(dyn, self.X, [0, dt], args=(U, self.P, ))
        else:
            ret = self.c_impl.run(self.X, dt, U)
        

        min_speed = 1.; max_phi = mu.rad_of_deg(60); max_theta = mu.rad_of_deg(60); 
        if self.X[s_v] < min_speed or math.fabs(self.X[s_phi]) > max_phi or  math.fabs(self.X[s_theta]) > max_theta:
            self.status =  DynamicModel.sta_crashed

        self.emit("update_fast")
        if self.update_cnt % 5 == 0: self.emit("update_slow")
        self.update_cnt += 1
        
    def name(self): return self.P.name

    def jacobian(self):
        A,B = mu.num_jacobian(self.X, self.U, self.P, dyn)
        return A, B



class DynamicModelJSBsim(DynamicModel):
    
    def __init__(self, X0=None):
        DynamicModel.__init__(self)
        try:
            self.jsbsim = __import__("JSBSim")
            self.fdex = self.jsbsim.FGFDMExec( )
            self.fdex.SetDebugLevel( 0 )
            self.fdex.SetAircraftPath('/usr/share/games/flightgear/Aircraft')
            self.fdex.LoadModel('c172p')
            propulsion = self.fdex.GetPropulsion( )
            fcs = self.fdex.GetFCS( )
            fcs.SetThrottleCmd( 0, 0.25 )
            fcs.SetMixtureCmd( 0, 0.87 )
            propulsion.SetMagnetos( 3 )
            propulsion.SetStarter( 1 )
        except ImportError:
            self.c_impl = None
            print "Error: could not load native JSBsim  library"
    
    def reset(self):
        DynamicModel.reset(self)
        try:
            fdex = self.fdex
        except AttributeError:
            pass
        else:
            ic = self.fdex.GetIC( )
            ic.SetLatitudeDegIC( 47.0 ) ; ic.SetLongitudeDegIC( 122.0 ); ic.SetAltitudeFtIC( 1500. )
            ic.SetUBodyFpsIC( 70.0 ); ic.SetVBodyFpsIC( 0.0 ) ;ic.SetWBodyFpsIC( 0.0 )
            ic.SetPhiDegIC( 0.0 ); ic.SetThetaDegIC( 0.0 ); ic.SetPsiDegIC( 0.0 )
            self.fdex.RunIC( )
            self.fdex.DoTrim( self.jsbsim.tGround )


    def run(self, U, dt):
        fcs = self.fdex.GetFCS( )
        fcs.SetDaCmd( U[1+i_da1] ) # aileron
        fcs.SetDeCmd( U[1+i_de1] ) # elevator
        fcs.SetDrCmd( U[1+i_dr1] ) # rudder
        fcs.SetThrottleCmd( 0, U[0] )

        self.fdex.Run()
        propagate = self.fdex.GetPropagate()
        self.X[s_p], self.X[s_q], self.X[s_r] = propagate.GetPQR()
        self.X[s_phi], self.X[s_theta], self.X[s_psi] = propagate.GetEuler( )
        u,v,w = propagate.GetUVW( )
        self.X[s_v] = math.sqrt(u**2+v**2+w**2)
#        location = propagate.GetLocation( )
#        ecx, ecy, ecz = location.GetEntries( )
#        self.X[s_x] = ecx
#        self.X[s_y] = ecy
#        self.X[s_z] = ecz
        self.X[s_x] = propagate.GetLatitude()
        self.X[s_y] = propagate.GetLongitude()
        self.X[s_z] = propagate.Geth()
        #        
        self.emit("update_fast")
        if self.update_cnt % 5 == 0: self.emit("update_slow")
        self.update_cnt += 1

        
    def name(self): return "JSBsim"

    
